from typing import Dict, Any, Optional
from agentscope.memory import Mem0LongTermMemory
from app.schemas.memory import MemoryConfig
from app.core.config import settings
from app.core.models import embedding_model, chat_model
from mem0.vector_stores.configs import VectorStoreConfig

class MemoryService:
    _instances: Dict[str, Mem0LongTermMemory] = {}

    @classmethod
    async def get_memory_instance(
        cls,
        user_id: str,
        config: Optional[MemoryConfig] = None
    ) -> Mem0LongTermMemory:
        """获取或创建内存实例"""
        if user_id not in cls._instances:
            if not config:
                config = MemoryConfig(user_name=user_id)
            
            # 创建 Milvus 向量存储配置
            vector_store_config = VectorStoreConfig(
                provider="milvus",
                config={
                    "host": settings.MILVUS_HOST,
                    "port": settings.MILVUS_PORT,
                    "user": settings.MILVUS_USER,
                    "password": settings.MILVUS_PASSWORD,
                    "collection_name": f"{settings.MILVUS_COLLECTION}_{user_id}",
                    "dimension": settings.MILVUS_DIMENSION,
                    "index_params": {
                        "metric_type": "L2",
                        "index_type": "IVF_FLAT",
                        "params": {"nlist": 1024}
                    }
                }
            )
            
            cls._instances[user_id] = Mem0LongTermMemory(
                agent_name=config.agent_name,
                user_name=config.user_name,
                run_name=config.run_name,
                default_memory_type=config.default_memory_type,
                vector_store_config=vector_store_config,
                model=chat_model,
                embedding_model=embedding_model
            )
        
        return cls._instances[user_id]

    @classmethod
    async def record_memory(
        cls,
        user_id: str,
        thinking: str,
        content: list[str],
        memory_type: Optional[str] = None,
        infer: bool = True,
        agent_name: Optional[str] = None,
        run_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """记录记忆到向量数据库"""
        memory_instance = await cls.get_memory_instance(
            user_id,
            MemoryConfig(
                agent_name=agent_name,
                user_name=user_id,
                run_name=run_name,
                default_memory_type=memory_type
            )
        )
        
        result = await memory_instance.record_to_memory(
            thinking=thinking,
            content=content,
            memory_type=memory_type,
            infer=infer
        )

        return {
            "message": "Memory recorded successfully",
            "details": result.content[0].text if result.content else None
        }

    @classmethod
    async def retrieve_memory(
        cls,
        user_id: str,
        keywords: list[str],
        limit: int = 5,
        memory_type: Optional[str] = None,
        agent_name: Optional[str] = None,
        run_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """从向量数据库检索记忆"""
        memory_instance = await cls.get_memory_instance(
            user_id,
            MemoryConfig(
                agent_name=agent_name,
                user_name=user_id,
                run_name=run_name,
                default_memory_type=memory_type
            )
        )
        
        result = await memory_instance.retrieve_from_memory(
            keywords=keywords,
            limit=limit
        )
        
        memories = result.content[0].text.split("\n") if result.content else []
        
        return {
            "memories": memories,
            "total_count": len(memories)
        }